from langchain import hub
from langchain.agents import tool, AgentExecutor, XMLAgent
from langchain.agents.format_scratchpad import format_xml
from langchain.agents.output_parsers import XMLAgentOutputParser
from langchain.chains import LLMChain
from langchain.chat_models import ChatAnthropic
from langchain.tools.render import render_text_description


@tool
def search(query: str) -> str:
    """Search things about current events."""
    return "32 degrees"


tools = [search]
model = ChatAnthropic(model="claude-2")


def use_expression_create():
    prompt = hub.pull("hwchase17/xml-agent")

    prompt = prompt.partial(
        tools=render_text_description(tools),
        tool_names=", ".join([t.name for t in tools]),
    )

    llm_with_stop = model.bind(stop=["</tool_input>"])

    agent = (
            {
                "question": lambda x: x["question"],
                "agent_scratchpad": lambda x: format_xml(x["intermediate_steps"]),
            }
            | prompt
            | llm_with_stop
            | XMLAgentOutputParser()
    )
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    print(agent_executor.invoke({"question": "whats the weather in New york?"}))


def use_self_agent():
    # 使用现成的代理
    chain = LLMChain(
        llm=model,
        prompt=XMLAgent.get_default_prompt(),
        output_parser=XMLAgent.get_default_output_parser(),
    )
    agent = XMLAgent(tools=tools, llm_chain=chain)

    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    agent_executor.invoke({"input": "whats the weather in New york?"})


if __name__ == '__main__':
    use_self_agent()